#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/epoll.h>
#include <unistd.h>

#define MAX_EVENTS 1024
#define BUFFER_SIZE 1024
#define PORT 4000

void set_nonblocking(int sockfd)
{
    int opts = fcntl(sockfd, F_GETFL);
    if (opts < 0) {
        perror("fcntl(F_GETFL)");
        exit(EXIT_FAILURE);
    }
    opts = opts | O_NONBLOCK;
    if (fcntl(sockfd, F_SETFL, opts) < 0) {
        perror("fcntl(F_SETFL)");
        exit(EXIT_FAILURE);
    }
}

int main()
{
    int udp_fd, epoll_fd;
    struct sockaddr_in serv_addr, cli_addr;
    struct epoll_event ev, events[MAX_EVENTS];
    char buffer[BUFFER_SIZE];
    socklen_t addr_len = sizeof(cli_addr);

    // 创建 UDP socket
    if ((udp_fd = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
        perror("socket");
        exit(EXIT_FAILURE);
    }

    set_nonblocking(udp_fd);

    // 设置服务器地址结构
    memset(&serv_addr, 0, sizeof(serv_addr));
    serv_addr.sin_family = AF_INET;
    serv_addr.sin_addr.s_addr = htonl(INADDR_ANY);
    serv_addr.sin_port = htons(PORT);

    // 绑定 socket
    if (bind(udp_fd, (struct sockaddr*)&serv_addr, sizeof(serv_addr)) < 0) {
        perror("bind");
        exit(EXIT_FAILURE);
    }

    // 创建 epoll 实例
    if ((epoll_fd = epoll_create1(0)) < 0) {
        perror("epoll_create1");
        exit(EXIT_FAILURE);
    }

    // 将 UDP socket 加入 epoll 监听列表
    ev.events = EPOLLIN;
    ev.data.fd = udp_fd;
    if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, udp_fd, &ev) < 0) {
        perror("epoll_ctl: udp_fd");
        exit(EXIT_FAILURE);
    }

    while (1) {
        int nfds = epoll_wait(epoll_fd, events, MAX_EVENTS, -1);
        if (nfds < 0) {
            perror("epoll_wait");
            exit(EXIT_FAILURE);
        }

        for (int i = 0; i < nfds; i++) {
            if (events[i].events & EPOLLIN) {
                // 接收数据
                int sock_fd = events[i].data.fd;
                ssize_t n = recvfrom(sock_fd, buffer, BUFFER_SIZE, 0, (struct sockaddr*)&cli_addr, &addr_len);
                if (n > 0) {
                    // 回显数据
                    sendto(sock_fd, buffer, n, 0, (struct sockaddr*)&cli_addr, addr_len);
                } else if (n < 0 && errno != EAGAIN && errno != EWOULDBLOCK) {
                    perror("recvfrom");
                }
            }
        }
    }

    close(udp_fd);
    close(epoll_fd);
    return 0;
}
